import numpy as np
from embed import get_embedding
from store import EntertainmentRAGStore
from loguru import logger
from typing import List, Dict, Any


def search_similar(query: str, top_k: int = 3) -> List[Dict]:
    store = EntertainmentRAGStore()
    if len(store) == 0:
        return []

    # 关键：设置搜索时查看的聚类数量（nprobe越大，精度越高但速度越慢）
    # 建议值：5-100（根据nlist调整，通常为nlist的1/10）
    store.index.nprobe = 10

    query_vec = np.array([get_embedding(query)], dtype=np.float32)
    distances, indices = store.index.search(query_vec, top_k)

    results = []
    for i, idx in enumerate(indices[0]):
        if idx < len(store.metadata):
            results.append({
                "text": store.metadata[idx]["text"],
                "score": float(distances[0][i])  # 距离越小越相似
            })
    return results


if __name__ == '__main__':
    results = search_similar('aaaaabbbbb', 3)
    logger.info(f"搜索结果：{results}")